import scipy.stats as stats
import matplotlib.pyplot as plt
import numpy as np
import matplotlib

matplotlib.use('TkAgg')
# 设置中文字体
matplotlib.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
matplotlib.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题


def calculate_admission_probability(enrollment_plan, candidate_rank,
                                    historical_min_rank_list, historical_enrollment_list,
                                    heat_change_coefficient=0):
    """
    计算考生被专业录取的概率

    :param enrollment_plan: 目标高校专业A在上海的招生计划人数
    :param candidate_rank: 考生个人的总排名
    :param historical_min_rank_list: 历年录取最低位次数据列表
    :param historical_enrollment_list: 历年录取人数数据列表
    :param heat_change_coefficient: 专业热度变化调整系数，默认为0
    :return: 考生被专业录取的概率
    """
    # 计算历年录取最低位次的统计量
    # 计算平均最低位次
    average_min_rank = sum(historical_min_rank_list) / len(historical_min_rank_list)
    # 计算标准差
    variance = sum((x - average_min_rank) ** 2 for x in historical_min_rank_list) / len(historical_min_rank_list)
    standard_deviation = variance ** 0.5

    print(f"历年录取最低位次的标准差为: {standard_deviation}")

    # 构建概率计算模型（基于正态分布假设）
    # 计算标准分数（Z - score）
    if standard_deviation == 0:
        # 处理标准差为 0 的情况
        distance = abs(candidate_rank - average_min_rank)
        if candidate_rank < average_min_rank:
            # 当考生排名小于平均最低位次时，距离越大z_score越小
            z_score = -distance / (max(historical_min_rank_list) - min(historical_min_rank_list) + 1)
        elif candidate_rank > average_min_rank:
            # 当考生排名大于平均最低位次时，距离越大z_score越大
            z_score = distance / (max(historical_min_rank_list) - min(historical_min_rank_list) + 1)
        else:
            z_score = 0
    else:
        z_score = (candidate_rank - average_min_rank) / standard_deviation

    # 计算录取概率
    admission_probability = stats.norm.sf(z_score)

    # 考虑招生计划调整概率
    # 计算历年平均录取人数
    average_historical_enrollment = sum(historical_enrollment_list) / len(historical_enrollment_list)
    # 调整因子
    adjustment_factor = enrollment_plan / average_historical_enrollment
    # 调整后的概率
    adjusted_admission_probability = admission_probability * adjustment_factor
    # 确保概率在[0, 1]范围内
    adjusted_admission_probability = max(0, min(adjusted_admission_probability, 1))

    # 考虑专业热度变化因素
    # 调整后的概率
    final_admission_probability = adjusted_admission_probability * (1 + heat_change_coefficient)
    # 确保最终概率在[0, 1]范围内
    final_admission_probability = max(0, min(final_admission_probability, 1))

    return final_admission_probability


# 目标高校专业A在上海的招生计划人数
enrollment_plan = 30
# 假设获取到的历年录取最低位次数据（这里以5年为例），存储在列表中
historical_min_rank_list = [38393,42365]
# 假设获取到的历年录取人数数据（这里以5年为例），存储在列表中
historical_enrollment_list = [24,38]
# 假设评估得到的热度变化调整系数h（这里假设h = 0）
heat_change_coefficient = 0

# 模拟考生排名，仅保留一个排名数据
ranks = [35455]
probabilities = []
for rank in ranks:
    prob = calculate_admission_probability(enrollment_plan, rank,
                                           historical_min_rank_list, historical_enrollment_list,
                                           heat_change_coefficient)
    probabilities.append(prob)

# 计算特定排名的录取概率
specific_ranks = ranks
for rank in specific_ranks:
    rank_prob = calculate_admission_probability(enrollment_plan, rank,
                                                historical_min_rank_list, historical_enrollment_list,
                                                heat_change_coefficient)
    print(f"排名{rank}的录取概率为: {rank_prob * 100:.2f}%")

# 由于只有一个数据点，不适合绘制曲线，这里改为打印数据点
if len(ranks) == 1:
    print(f"数据点: ({ranks[0]}, {probabilities[0]:.4f})")
else:
    # 绘制曲线
    plt.plot(ranks, probabilities)
    plt.xlabel('考生排名')
    plt.ylabel('录取概率')
    plt.title('考生不同排名下该专业的录取概率曲线')

    # 标注关键点
    for rank in specific_ranks:
        rank_prob = calculate_admission_probability(enrollment_plan, rank,
                                                    historical_min_rank_list, historical_enrollment_list,
                                                    heat_change_coefficient)
        plt.annotate(f'({rank}, {rank_prob:.4f})', xy=(rank, rank_prob), xytext=(rank + 50, rank_prob + 0.05),
                     arrowprops=dict(facecolor='black', shrink=0.05))

    plt.show()
